"""
Conservation law monitoring and energy function computation

This module provides trackers for monitoring conservation laws, energy functionals,
and other invariants during PDE evolution.
"""

using ..Fields
using ..PDEs
using LinearAlgebra
using Statistics

"""
    ConservationTracker

Tracks conservation laws and invariants during PDE evolution.
"""
struct ConservationTracker
    quantities::Dict{Symbol, Function}  # name -> computation function
    history::Dict{Symbol, Vector{Float64}}
    tolerance::Float64
    check_frequency::Int
    current_step::Int
    
    function ConservationTracker(quantities::Dict{Symbol, Function}; 
                               tolerance=1e-10, check_frequency=1)
        history = Dict(name => Float64[] for name in keys(quantities))
        new(quantities, history, tolerance, check_frequency, 0)
    end
end

"""
    track!(tracker::ConservationTracker, pde::AbstractPDE, state, t)

Monitor conservation laws at the current time step.
"""
function track!(tracker::ConservationTracker, pde::AbstractPDE, state, t)
    tracker.current_step += 1
    
    if tracker.current_step % tracker.check_frequency == 0
        for (name, compute_func) in tracker.quantities
            value = compute_func(pde, state, t)
            push!(tracker.history[name], value)
            
            # Check for conservation violations
            if length(tracker.history[name]) > 1
                initial_value = tracker.history[name][1]
                relative_change = abs(value - initial_value) / abs(initial_value)
                
                if relative_change > tracker.tolerance
                    @warn "Conservation violation detected for $name: " *
                          "relative change = $relative_change"
                end
            end
        end
    end
end

"""
    get_conservation_errors(tracker::ConservationTracker)

Get relative errors in conserved quantities.
"""
function get_conservation_errors(tracker::ConservationTracker)
    errors = Dict{Symbol, Float64}()
    
    for (name, history) in tracker.history
        if length(history) > 1
            initial_value = history[1]
            final_value = history[end]
            relative_error = abs(final_value - initial_value) / abs(initial_value)
            errors[name] = relative_error
        end
    end
    
    return errors
end

"""
    EnergyTracker

Specialized tracker for energy functionals and Lyapunov functions.
"""
struct EnergyTracker
    energy_function::Function
    gradient_function::Union{Function, Nothing}
    history::Vector{Float64}
    gradient_history::Vector{Float64}
    dissipation_rate::Vector{Float64}
    check_frequency::Int
    current_step::Int
    
    function EnergyTracker(energy_func::Function; 
                          gradient_func=nothing, check_frequency=1)
        new(energy_func, gradient_func, Float64[], Float64[], Float64[],
            check_frequency, 0)
    end
end

"""
    track!(tracker::EnergyTracker, pde::AbstractPDE, state, t)

Track energy evolution and dissipation.
"""
function track!(tracker::EnergyTracker, pde::AbstractPDE, state, t)
    tracker.current_step += 1
    
    if tracker.current_step % tracker.check_frequency == 0
        # Compute energy
        energy = tracker.energy_function(pde, state, t)
        push!(tracker.history, energy)
        
        # Compute energy gradient if available
        if tracker.gradient_function !== nothing
            gradient = tracker.gradient_function(pde, state, t)
            gradient_norm = norm(gradient)
            push!(tracker.gradient_history, gradient_norm)
        end
        
        # Estimate dissipation rate
        if length(tracker.history) > 1
            dt = 0.01  # Should get from solver
            energy_change = tracker.history[end] - tracker.history[end-1]
            dissipation = -energy_change / dt
            push!(tracker.dissipation_rate, dissipation)
        end
    end
end

"""
    is_energy_decreasing(tracker::EnergyTracker; window_size=10)

Check if energy is consistently decreasing (for gradient flows).
"""
function is_energy_decreasing(tracker::EnergyTracker; window_size=10)
    if length(tracker.history) < window_size
        return true  # Not enough data
    end
    
    recent_history = tracker.history[end-window_size+1:end]
    
    # Check if energy is generally decreasing
    decreasing_count = 0
    for i in 2:length(recent_history)
        if recent_history[i] <= recent_history[i-1]
            decreasing_count += 1
        end
    end
    
    return decreasing_count / (window_size - 1) > 0.8
end

"""
    compute_mass(pde::AbstractPDE, state, t=0.0)

Compute total mass: M = ∫ u dV
"""
function compute_mass(pde::AbstractPDE, state, t=0.0)
    if isa(state, ScalarField)
        return sum(data(state))
    elseif isa(state, FieldCollection)
        return sum(sum(data(field)) for field in state)
    else
        return 0.0
    end
end

"""
    compute_l2_norm(pde::AbstractPDE, state, t=0.0)

Compute L² norm: ||u||² = ∫ u² dV
"""
function compute_l2_norm(pde::AbstractPDE, state, t=0.0)
    if isa(state, ScalarField)
        return sum(data(state).^2)
    elseif isa(state, FieldCollection)
        return sum(sum(data(field).^2) for field in state)
    else
        return 0.0
    end
end

"""
    compute_allen_cahn_energy(pde::AllenCahnPDE, state, t=0.0)

Compute Allen-Cahn free energy: E = ∫[f(c) + (γ/2)|∇c|²] dV
"""
function compute_allen_cahn_energy(pde::AllenCahnPDE, state, t=0.0)
    return free_energy(pde, state)
end

"""
    compute_cahn_hilliard_energy(pde::CahnHilliardPDE, state, t=0.0)

Compute Cahn-Hilliard free energy.
"""
function compute_cahn_hilliard_energy(pde::CahnHilliardPDE, state, t=0.0)
    return free_energy(pde, state)
end

"""
    compute_wave_energy(pde::WavePDE, state, t=0.0)

Compute wave energy: E = ∫[v²/2 + c²|∇u|²/2] dV
"""
function compute_wave_energy(pde::WavePDE, state, t=0.0)
    return wave_energy(pde, state)
end

"""
    compute_ks_energy(pde::KuramotoSivashinskyPDE, state, t=0.0)

Compute Kuramoto-Sivashinsky energy.
"""
function compute_ks_energy(pde::KuramotoSivashinskyPDE, state, t=0.0)
    return get_energy(pde, state)
end

"""
    compute_swift_hohenberg_energy(pde::SwiftHohenbergPDE, state, t=0.0)

Compute Swift-Hohenberg energy functional.
"""
function compute_swift_hohenberg_energy(pde::SwiftHohenbergPDE, state, t=0.0)
    return free_energy(pde, state)
end

"""
    create_conservation_tracker(pde::AbstractPDE; tolerance=1e-10)

Create appropriate conservation tracker for the PDE type.
"""
function create_conservation_tracker(pde::AbstractPDE; tolerance=1e-10)
    quantities = Dict{Symbol, Function}()
    
    # Always track basic quantities
    quantities[:mass] = compute_mass
    quantities[:l2_norm] = compute_l2_norm
    
    # Add PDE-specific conserved quantities
    if isa(pde, CahnHilliardPDE)
        # Mass is conserved in Cahn-Hilliard
        quantities[:energy] = compute_cahn_hilliard_energy
        
    elseif isa(pde, WavePDE)
        # Energy is conserved in wave equation
        quantities[:energy] = compute_wave_energy
        
    elseif isa(pde, AllenCahnPDE)
        # Energy decreases in Allen-Cahn (not conserved)
        quantities[:energy] = compute_allen_cahn_energy
        
    elseif isa(pde, SwiftHohenbergPDE)
        quantities[:energy] = compute_swift_hohenberg_energy
        
    elseif isa(pde, KuramotoSivashinskyPDE)
        quantities[:energy] = compute_ks_energy
    end
    
    return ConservationTracker(quantities, tolerance=tolerance)
end

"""
    create_energy_tracker(pde::AbstractPDE)

Create energy tracker for the specific PDE type.
"""
function create_energy_tracker(pde::AbstractPDE)
    if isa(pde, AllenCahnPDE)
        return EnergyTracker(compute_allen_cahn_energy)
    elseif isa(pde, CahnHilliardPDE)
        return EnergyTracker(compute_cahn_hilliard_energy)
    elseif isa(pde, WavePDE)
        return EnergyTracker(compute_wave_energy)
    elseif isa(pde, SwiftHohenbergPDE)
        return EnergyTracker(compute_swift_hohenberg_energy)
    elseif isa(pde, KuramotoSivashinskyPDE)
        return EnergyTracker(compute_ks_energy)
    else
        # Default: track L² norm as "energy"
        return EnergyTracker(compute_l2_norm)
    end
end

"""
    LyapunovTracker

Track Lyapunov functional for stability analysis.
"""
struct LyapunovTracker
    lyapunov_function::Function
    history::Vector{Float64}
    is_decreasing::Bool
    violation_count::Int
    
    function LyapunovTracker(lyapunov_func::Function)
        new(lyapunov_func, Float64[], true, 0)
    end
end

"""
    track!(tracker::LyapunovTracker, pde::AbstractPDE, state, t)

Track Lyapunov functional evolution.
"""
function track!(tracker::LyapunovTracker, pde::AbstractPDE, state, t)
    value = tracker.lyapunov_function(pde, state, t)
    push!(tracker.history, value)
    
    # Check if Lyapunov property is maintained
    if length(tracker.history) > 1
        if value > tracker.history[end-1]
            tracker.is_decreasing = false
            tracker.violation_count += 1
        end
    end
end

"""
    InstabilityTracker

Track indicators of instability development.
"""
struct InstabilityTracker
    indicators::Dict{Symbol, Vector{Float64}}
    thresholds::Dict{Symbol, Float64}
    warnings::Vector{String}
    
    function InstabilityTracker()
        indicators = Dict{Symbol, Vector{Float64}}(
            :max_value => Float64[],
            :gradient_norm => Float64[],
            :energy_growth => Float64[]
        )
        thresholds = Dict{Symbol, Float64}(
            :max_value => 1e6,
            :gradient_norm => 1e6,
            :energy_growth => 1e6
        )
        new(indicators, thresholds, String[])
    end
end

"""
    track!(tracker::InstabilityTracker, pde::AbstractPDE, state, t)

Track instability indicators.
"""
function track!(tracker::InstabilityTracker, pde::AbstractPDE, state, t)
    # Maximum field value
    if isa(state, ScalarField)
        max_val = maximum(abs.(data(state)))
    elseif isa(state, FieldCollection)
        max_val = maximum(maximum(abs.(data(field))) for field in state)
    else
        max_val = 0.0
    end
    push!(tracker.indicators[:max_value], max_val)
    
    # Gradient magnitude
    if isa(state, ScalarField)
        grad = gradient(state)
        if isa(grad, VectorField)
            grad_norm = maximum(norm.(data(grad)))
        else
            grad_norm = maximum(abs.(grad))
        end
    else
        grad_norm = 0.0
    end
    push!(tracker.indicators[:gradient_norm], grad_norm)
    
    # Energy growth rate
    energy = compute_l2_norm(pde, state, t)
    if length(tracker.indicators[:energy_growth]) > 0
        growth = energy - tracker.indicators[:energy_growth][end]
        push!(tracker.indicators[:energy_growth], growth)
    else
        push!(tracker.indicators[:energy_growth], 0.0)
    end
    
    # Check thresholds
    for (name, values) in tracker.indicators
        if !isempty(values) && values[end] > tracker.thresholds[name]
            warning_msg = "Instability detected: $name = $(values[end])"
            if warning_msg ∉ tracker.warnings
                push!(tracker.warnings, warning_msg)
                @warn warning_msg
            end
        end
    end
end

"""
    get_instability_score(tracker::InstabilityTracker)

Compute overall instability score (0 = stable, 1 = unstable).
"""
function get_instability_score(tracker::InstabilityTracker)
    score = 0.0
    count = 0
    
    for (name, values) in tracker.indicators
        if !isempty(values)
            normalized_value = values[end] / tracker.thresholds[name]
            score += min(normalized_value, 1.0)
            count += 1
        end
    end
    
    return count > 0 ? score / count : 0.0
end

# Export conservation and energy tracking
export ConservationTracker, EnergyTracker, LyapunovTracker, InstabilityTracker
export create_conservation_tracker, create_energy_tracker
export get_conservation_errors, is_energy_decreasing, get_instability_score
export compute_mass, compute_l2_norm